from torch.utils.data import Dataset
from PIL import Image
import os


class MyData(Dataset):
    
    def __init__(self, root_dir, label_dir):
        # 初始化函数，为后面的getitem和next方法提供所需要的量
        self.root_dir = root_dir # 根目录的路径
        self.label_dir = label_dir # label目录的路径，因为下一行要合起来，并且label名是文件夹名，所以这里的label_dir可以直接取对应的label，如："ants"
        self.path = os.path.join(root_dir, label_dir) # 将两个路径连起来
        self.img_path = os.listdir(self.path) # 获取该路径下所有文件的路径列表
        
    def __getitem__(self, idx):
        """获取数据集中的每一个图片，输入索引，得到对应的图片"""
        # idx是index索引的缩写
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.path, img_name)
        image = Image.open(img_item_path)
        label = self.label_dir
        return image, label
    
    def __len__(self):
        return len(self.img_path)


if __name__ == '__main__':
    root_dir = "./dataset/hymenoptera_data/train"
    ants_label_dir = "ants"
    bees_label_dir = "bees"
    ants_dataset = MyData(root_dir, ants_label_dir)
    bees_dataset = MyData(root_dir, bees_label_dir)
    img, label = ants_dataset[0]
    print(img, label)
    img.show()
    
